# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: leon

import ujson as json
import functools

from coin.base.book.types import TradeSide, create_market_trade_event
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.gopax.book_builder import BookBuilder
from coin.exchange.gopax.kr_rest.product import GopaxProduct


def to_trade_side(is_buy):
  if is_buy == 1:
    return TradeSide.TRADE_BUY_SIDE
  elif is_buy == 0:
    return TradeSide.TRADE_SELL_SIDE
  else:
    raise ValueError(is_buy)


class GopaxHandler(WebsocketHandler):
  def __init__(self, *, event_dist=None, feed_checker_config=None):
    super().__init__(event_dist=event_dist, feed_checker_config=feed_checker_config)
    self._assets_data = None
    self._trading_pairs_data = None

  def on_book(self, record, queue_data, topic_data):
    data = record.data

    if not isinstance(data, dict):
      data = json.loads(data.tobytes())

    event = data['n']
    if event == 'GetAssets':
      # They encode again.
      self._assets_data = json.loads(data['o'])['data']
      return
    elif event == 'GetTradingPairs':
      self._trading_pairs_data = json.loads(data['o'])['data']
      return
    else:
      if not all((self._assets_data, self._trading_pairs_data)):
        return

    o = json.loads(data['o'])
    if event == 'GetOrderBook':
      pair_id = data['pairId']
    else:
      pair_id = o['pairId']
    product = self.get_product_from_pair_id(pair_id)
    native_symbol = product.native_symbol
    if not self.is_subscribed('book', native_symbol):
      return

    if event == 'GetOrderBook':
      book = BookBuilder(native_symbol, book_builder_observer=self._book_builder_observer)
      bids = o['bid']
      asks = o['ask']
      book.snapshot(bids, asks)
      self._books[native_symbol] = book
    elif event == 'OrderBookEvent':
      if native_symbol not in self._books:
        return
      book = self._books[native_symbol]
      bids = o['data']['bid']
      asks = o['data']['ask']
      book.update(bids, asks)
    else:
      raise ValueError(data)
    book.update_timestamp(record.timestamp)
    book.update_topic_id(record.topic_id)
    self.ready = True

    if self._feed_checker:
      self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
    self.publish_event('book', book.symbol, book)

  def on_trade(self, record, queue_data, topic_data):
    data = record.data

    if not isinstance(data, dict):
      data = json.loads(data.tobytes())

    if not all((self._assets_data, self._trading_pairs_data)):
      return
    event = data['n']
    assert event == 'TradeEvent'
    o = json.loads(data['o'])
    pair_id = o['pairId']
    product = self.get_product_from_pair_id(pair_id)
    native_symbol = product.native_symbol
    if not self.is_subscribed('trade', native_symbol):
      return

    price = float(o['price'])
    qty = float(o['baseAmount'])
    side = to_trade_side(int(o['isBuy']))
    trade_event = create_market_trade_event(record.timestamp,
                                            native_symbol,
                                            price,
                                            qty,
                                            side,
                                            topic_id=record.topic_id)
    if self._feed_checker:
      self._feed_checker.on_trade(product=product,
                                  trade_event=trade_event,
                                  timestamp=record.timestamp)
    self.publish_event('trade', native_symbol, trade_event)

  @functools.lru_cache(maxsize=64)
  def get_product_from_pair_id(self, pair_id):
    quote_asset_id = None
    base_asset_id = None
    for pair in self._trading_pairs_data:
      if pair_id == pair['tradingPairId']:
        base_asset_id = pair['baseAssetId']
        quote_asset_id = pair['quoteAssetId']
        break
    assert all((quote_asset_id, base_asset_id))
    quote = None
    base = None
    for asset in self._assets_data:
      if quote_asset_id == asset['assetId']:
        quote = asset['isoAlpha3']
      elif base_asset_id == asset['assetId']:
        base = asset['isoAlpha3']
      if all((quote, base)):
        break
    assert all((quote, base))
    return GopaxProduct.FromStrNativeProduct('%s-%s' % (base, quote))
